import cv2
from ultralytics import YOLO
from deep_sort_realtime.deepsort_tracker import DeepSort
import numpy as np

# # 加载YOLOv8模型
# model = YOLO(r"yolov8n.pt")
#
# # 加载DeepSORT跟踪器
# deepsort_tracker = DeepSort(max_age=60, n_init=2, max_cosine_distance=0.5)
#
# # 打开视频流
# cap = cv2.VideoCapture("2.mp4")
# track_history = {}
#
# while cap.isOpened():
#     ret, frame = cap.read()
#     if not ret:
#         break
#
#     height, width = frame.shape[:2]
#     imgsz = min(1024, width, height)  # 限制输入尺寸
#
#     # 目标检测（使用动态输入尺寸，关闭自动缩放填充）
#     results = model(frame, imgsz=imgsz, conf=0.4)
#
#     detections = []
#     for result in results:
#         for box in result.boxes:
#             class_id = int(box.cls[0])
#             if class_id not in [0, 2]:  # 仅保留行人（0）和车辆（2）
#                 continue
#
#             # 获取原始坐标（已适配imgsz，无需手动映射）
#             x1, y1, x2, y2 = map(int, box.xyxy[0])
#
#             # 过滤无效检测框（宽高为负或超出图像范围）
#             if x2 <= x1 or y2 <= y1 or x1 < 0 or y1 < 0 or x2 > width or y2 > height:
#                 continue
#
#             score = box.conf[0].item()
#             detections.append(([x1, y1, x2, y2], score, class_id))
#
#     # DeepSORT轨迹跟踪
#     tracks = deepsort_tracker.update_tracks(detections, frame=frame)
#     for track in tracks:
#         if not track.is_confirmed():
#             continue
#
#         track_id = track.track_id
#         ltrb = track.to_ltrb()  # 获取跟踪后的坐标（左、上、右、下）
#
#         # 坐标平滑
#         if track_id not in track_history:
#             track_history[track_id] = [ltrb]
#         else:
#             track_history[track_id].append(ltrb)
#             if len(track_history[track_id]) > 5:
#                 track_history[track_id].pop(0)
#         smooth_ltrb = np.mean(track_history[track_id], axis=0).astype(int)
#         x1, y1, x2, y2 = smooth_ltrb
#
#         class_id = track.get_det_class()
#         color = (0, 255, 0) if class_id == 0 else (255, 0, 0)
#         label = f'P ID {track_id}' if class_id == 0 else f'V ID {track_id}'
#
#         # 绘制检测框和标签
#         cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
#         cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
#
#         # 标记违禁行为（示例：红色边框）
#         cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 1, cv2.LINE_AA)
#
#     cv2.imshow('Highway Detection', frame)
#     if cv2.waitKey(1) & 0xFF == ord('q'):
#         break
#
# cap.release()
# cv2.destroyAllWindows()

import cv2
from ultralytics import YOLO

# 加载预训练的 YOLOv8 模型
model = YOLO('yolo11n.pt')

# 打开视频流，可以替换为具体的视频文件路径或摄像头设备编号
cap = cv2.VideoCapture('2.mp4')

# 定义检测类别：0=行人，2=车辆（根据YOLOv8标签定义，需确认实际类别ID）
DETECT_CLASSES = [0, 2]
COLORS = {0: (0, 255, 0), 2: (255, 0, 0)}  # 行人绿色，车辆红色

while True:
    # 读取视频帧
    ret, frame = cap.read()
    if not ret:
        break

    # 使用 YOLOv8 模型进行检测，指定检测类别
    results = model.track(frame, classes=DETECT_CLASSES)

    # 遍历检测结果
    for result in results:
        boxes = result.boxes
        for box in boxes:
            # 获取类别ID和坐标
            cls_id = int(box.cls[0])
            x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
            label = model.names[cls_id]
            # 获取颜色和标签文本
            color = COLORS.get(cls_id, (255, 255, 0))  # 默认黄色
            # 在图像上绘制检测框和标签
            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)
            cv2.putText(frame, label, (x1, y1 - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)

    # 显示处理后的帧
    cv2.imshow('Detection', frame)

    # 按 'q' 键退出循环
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 释放资源
cap.release()
cv2.destroyAllWindows()
